#!/usr/bin/python

""" 
    This is the code to accompany the Lesson 3 (decision tree) mini-project.

    Use a Decision Tree to identify emails from the Enron corpus by author:    
    Sara has label 0
    Chris has label 1
"""
    
import sys
from time import time
sys.path.append("../tools/")
from email_preprocess import preprocess


### features_train and features_test are the features for the training
### and testing datasets, respectively
### labels_train and labels_test are the corresponding item labels
features_train, features_test, labels_train, labels_test = preprocess()




#########################################################
### your code goes here ###
from sklearn.tree import DecisionTreeClassifier as DTC
import numpy as np

if 0:
  hypotheses = []
  hypothesis_weights = []
  
  N = len(features_train)
  d = np.ones(N)/N
  num_iterations = 3
  
  labels_train = np.array(map(lambda x: (x==0) and -1 or 1, labels_train))
  print(labels_train[:20])
  for t in range(num_iterations):
    h = DTC(max_depth=1)
  
    h.fit(features_train, labels_train, sample_weight=d)
    pred = h.predict(features_train)
  
    eps = d.dot(pred != labels_train)
    print("EPS is %s" %eps)
    alpha = (np.log(1-eps) - np.log(eps)) / 2
  
    d = d * np.exp(-alpha*labels_train*pred)
    d = d / d.sum()
  
    hypotheses.append(h)
    hypothesis_weights.append(alpha)
  
  pred = np.zeros(len(features_test))
  for (h, alpha) in zip(hypotheses, hypothesis_weights):
    pred = pred + alpha * h.predict(features_test)
  pred = np.sign(pred).astype(int).tolist()
  pred = map(lambda x: (x==-1) and 0 or 1, pred)

#if 0:
#  from sklearn.ensemble import 

if 1: # decision tree
  clf = DTC(min_samples_split=40)
  
  t0 = time()
  clf.fit(features_train, labels_train)
  print("It took %s s to train" % round(time()-t0,3))
  
  t0 = time()
  pred = clf.predict(features_test)
  print("It took %s s to predict" % round(time()-t0,3))

from sklearn.metrics import accuracy_score
accuracy = accuracy_score(labels_test, pred)
print("Accuracy is %s" % round(accuracy, 3))



#########################################################


